{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.integrate import odeint\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "m1 = 1\n",
    "m2 = 1\n",
    "L1 = 0.5\n",
    "L2 = 0.5\n",
    "g = 9.81\n",
    "\n",
    "def M(theta):\n",
    "    return np.array([[m1 * L1**2 + m2 * (L1**2 + 2 * L1 * L2 * math.cos(theta[0][1]) + L2**2), m2 * (L1 * L2 * math.cos(theta[0][1]) + L2**2)], \n",
    "                     [m2 *(L1 * L2 * math.cos(theta[0][1]) + L2**2)                          , m2 * L2**2]])\n",
    "\n",
    "\n",
    "def c(theta, theta_dot):\n",
    "    return np.array([[- m2 * L1 * L2 * math.sin(theta[0][1]) * (2 * theta_dot[0][0] * theta_dot[0][1] + theta_dot[0][1] ** 2)],\n",
    "                     [m2 * L1 * L2 * theta_dot[0][0] ** 2 * math.sin(theta[0][1])]])\n",
    "\n",
    "def g_f(theta):\n",
    "    return np.array([[(m1 + m2) * L1 * g * math.cos(theta[0][0]) + m2 * g * L2 * math.cos(theta[0][0] + theta[0][1])], \n",
    "                     [m2 * g * L2 * math.cos(theta[0][0] + theta[0][1])]])\n",
    "\n",
    "\n",
    "theta0 = np.array([[0, np.pi * 0.5]])\n",
    "# theta0 = np.array([[-np.pi * 0.5, 0]])\n",
    "# theta0 = np.array([[0, 0]])\n",
    "theta1 = np.array([[0, -np.pi * 0.5]])\n",
    "theta_dot0 = np.array([[0, 0]])\n",
    "torque = np.array([[0, 0]])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(60001, 2)\n"
     ]
    }
   ],
   "source": [
    "y_res = [theta0[0]]\n",
    "\n",
    "theta = theta0\n",
    "theta_dot = theta_dot0\n",
    "dT = 0.001\n",
    "T = 60\n",
    "t = np.arange(0, T, 0.001)\n",
    "N = t.size\n",
    "\n",
    "def model(y, t):\n",
    "    theta, theta_dot = np.array([y[0:2]]), np.array([y[2:]])\n",
    "    # print(theta.shape, theta_dot.shape)\n",
    "    theta_next, theta_dot_next = (theta.T + dT * theta_dot.T).T, (theta_dot.T + dT * np.dot(np.linalg.inv(M(theta)), (torque.T - g_f(theta) - c(theta, theta_dot)))).T\n",
    "    return np.hstack((theta_next.reshape(-1), theta_dot_next.reshape(-1)))\n",
    "\n",
    "# result = odeint(model, np.hstack((theta.reshape(-1), theta_dot.reshape(-1))), t)\n",
    "# print(result)\n",
    "for i in range(N):\n",
    "    theta_tmp = (theta.T + 0.001 * theta_dot.T).T\n",
    "    theta_dot = (theta_dot.T + 0.001 * np.dot(np.linalg.inv(M(theta)), (torque.T - g_f(theta) - c(theta, theta_dot)))).T\n",
    "    theta = theta_tmp\n",
    "    # print(theta)\n",
    "    y_res.append(theta[0])\n",
    "res = np.array(y_res)\n",
    "print(res.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(array([0., 1., 1.]), array([0., 0., 1.]))\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAQoAAAD8CAYAAACPd+p5AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAQ8ElEQVR4nO3dbWwd1Z3H8e8vxnYCCaSFpICBgtQou7TdboDlQdkX0G6bgFChiK5gpZZtkap2QWqlbrTQatt9t0iR+gKBoKyKKBKCIpGk0ZKSstCIsmpaHhIIIZimgS52UB0e8uDEOHHy3xdzbRzn3nsumfF9GP8+kuU7c0/mzCHxjzMzx/eviMDMrJ5ZrT4BM2t/DgozS3JQmFmSg8LMkhwUZpbkoDCzpNxBIelsSb+RtE3SVknfrdJGku6UtF3Sy5IuyNuvmTXPCQUcYwz4fkS8KGke8IKkJyPi1UltrgQWVb4uAe6pfDezDpB7RhERb0fEi5XX+4BtQN+UZtcAD0ZmIzBf0hl5+zaz5ihiRjFB0rnAEuD3U97qA96atD1Q2fd2lWN8C/gWwOzZsy8855xzijzFtnDkyBFmzSrf7aGyjgvKO7bXX3/9nYhYkGpXWFBImgs8BnwvIvZOfbvKH6m6djwi7gPuA1i8eHH09/cXdYptY8OGDVx++eWtPo3ClXVcUN6xSfpzI+0KiUhJ3WQh8VBErKrSZAA4e9L2WcDOIvo2s+lXxFMPAT8DtkXET2o0Wwt8vfL041JgT0Qcc9lhZu2piEuPpcDXgC2SNlf2/QA4ByAi7gXWAVcB24EDwDcK6NfMmiR3UETEs1S/BzG5TQC35O3LzFqjfLdxzaxwDgozS3JQmFmSg8LMkhwUZpbkoDCzJAeFmSU5KMwsyUFhZkkOCjNLclCYWZKDwsySHBRmluSgMLMkB4WZJTkozCzJQWFmSQ4KM0sq6lO475c0JOmVGu9fLmmPpM2Vrx8V0a+ZNUdRdT0eAO4CHqzT5rcRcXVB/ZlZExUyo4iIZ4D3ijiWmbWfZt6juEzSS5J+JenTTezXzHIqtPZoHS8Cn4yIYUlXAWvIKpsfY3Lt0QULFrBhw4YmnWLzDA8Pe1wdpsxja4SykhsFHCgrUPzfEfGZBtq+CVwUEe/Ua+fao52lrOOC8o5N0gsRcVGqXVMuPSSdXik9iKSLK/2+24y+zSy/Qi49JD0MXA6cJmkA+DHQDRMlBa8HviNpDBgBboiipjJmNu0KCYqIuDHx/l1kj0/NrAN5ZaaZJTkozCzJQWFmSQ4KM0tyUJhZkoPCzJIcFGaW5KAwsyQHhZklOSjMLMlBYWZJDgozS3JQmFmSg8LMkhwUZpbkoDCzJAeFmSU5KMwsyUFhZknNqj0qSXdK2i7pZUkXFNGvmTVHUTOKB4Dldd6/kqzgzyKy4j73FNSvmTVBs2qPXgM8GJmNwHxJZxTRt5lNv2aVFOwD3pq0PVDZ9/bUhi4p2LnKOi4o99ga0aygUJV9VQsARcR9wH2QlRQsYxm3spanK+u4oNxja0SzgmIAOHvS9lnAzib1bXbc1mwaZOX6fgZ3j9C38WlWLFvMtUv6Wn1aTdesx6Nrga9Xnn5cCuyJiGMuO8zayZpNg9y26mUGd48AMLh7hNtXbWHNpsEWn1nzFfV49GHgd8BiSQOSbpb0bUnfrjRZB+wAtgP/BfxLEf2aTaeV6/v54NCRo/aNHDrMyvX9LTqj1mlW7dEAbimiL7Nm2VmZSTS6v8y8MtOshjPnz/lI+8vMQWFWw4pli+nuOvqB3ZzuLlYsW9yiM2odB4VZDdcu6eP6C8+a2O6bP4f/vO6zfuphZke74JyPAbD0zBP439s+PyNDAhwUZnVVXRU4AzkozOqpJMWsamuLZxAHhVkdR8JzCnBQmNU1HhPyjMLMavGEIuOgMKsjKnOKGT6hcFCY1TM+o3BQmFlNE1ceMzwpHBRm9YQvPcBBYVaXn3pkHBRmdRw54hkFOCjM6vLT0YyDwqwOP/XIOCjM6vA9ikxRn5m5XFJ/pWTgbVXev1zSHkmbK18/KqJfs+kWXpoJFPCZmZK6gLuBL5J9LP9zktZGxKtTmv42Iq7O259ZK8zwCUUhM4qLge0RsSMiDgKPkJUQNOt4vkeRKeJTuKuVC7ykSrvLJL1EVvjnXyNia7WDuaRg5yrjuLa/cQiAQ4cOlW5sH0URQdFIucAXgU9GxLCkq4A1ZJXNj/2DLinYsco4rtf0J+h/jZ6entKN7aMo4tIjWS4wIvZGxHDl9TqgW9JpBfRtNq0mLj1m+LVHEUHxHLBI0nmSeoAbyEoITpB0upT9p5Z0caXfdwvo22xa+dfMM7kvPSJiTNKtwHqgC7g/IraOlxOMiHuB64HvSBoDRoAbws+drAP4X2mmqJKC68jqi07ed++k13cBdxXRl1krzPQZhVdmmtUxPvH1PQozq8mXHhkHhVkdE7/r0dKzaD0HhVkdR3zpATgozOryEu6Mg8KsDt+iyDgozOrxpQfgoDCryzczMw4Kszr8eDTjoDCrY+J3PWb4lMJBYVaHZxQZB4VZHZWyHjP+B2Wmj9+srvFLj5l+N9NBYVaPcwJwUJjV5cejGQeFWR3hNdyAg8Ksrg9zYmYnhYPCrA4/Hc00q6SgJN1Zef9lSRcU0a+1hzWbBll6x9P88xP7WXrH06zZNNjqUyqMrzwyuYNiUknBK4HzgRslnT+l2ZVkdTwWkRX3uSdvv9Ye1mwa5PZVWxjcPQLA4O4Rbl+1pTRh4c+jyBTx4boTJQUBJI2XFJxce/Qa4MHKJ29vlDRf0hkR8XYB/VsLrVzfz8ihw0ftGzl0mB+s3kL/X/axcF4vC+b1snDebBZUXs/tLeQznZtqhudE00oKVmvTBxwTFC4p2FnGZxJTHTh4mHs2/Knqe71dcEqvmN8rTu7Jvo9vnzLp6+QeMavF/yt/a2AUgIMHR0vzd3Y8mlVSsJE22U6XFOwofRufrhoW8+d0c/Pfn8fQvlF27RtlaN8H7BoeZWjvKKNjRxg6EAwdqH+rcJbg1Lm9LJjby8KTp36fXZmpZLOUE3umZ5by9J5X4P/+TG9vb2n+zo5HEf91kyUFG2xjHWjFssXcvmrLUZcfc7q7+I8vf5prl/Qd0z4i2Dc6xtDeSQGyb5Rdw6Ps2js6ESa7hkd5b//B7L19o7yauEg9qacrC4+5vSyoFipzs0A59aQeZs1qbJayZtMgj70wkL3+40H+btNg1THNBEUExURJQWCQrKTgP01psxa4tXL/4hJgj+9PlMP4D87K9f0M7h6hb/4cVixbXPMHShInz+7m5NndfGrh3LrHPjh2hHf3j04KldFKqHwwESbj3/cfPMwb7+znjXf21z1m1yxx6kk9H4ZI5d7J0TOW2Wzc8S4/Xrt1IgD3j8Htq7YcNeaZpFklBdcBVwHbgQPAN/L2a+3j2iV9XLukr/BLqp4TZnHGKXM445Q5ddtFBHtHxqoGyNDe7JJnPGh2HzjEUOX1RzVy6DAr1/c7KI5XAyUFA7iliL7MppLEKSd2c8qJ3Xxq4by6bUfHDvPOcHZJMx4iU8Nl194P2Lnng6p/fmeNm7dl13nPqcxy6D2hi775c+ibX3+WsvSOpxjcfWxYnJn4c2XlJdxmVaxY9lfM6e46at+c7i5WLFvcojNqLc8ozKqYfJN25+4RPj5b/Ps1n52R9yfAQWFW0/hNWqisfZmhIQG+9DCzBjgozCzJQWFmSQ4KM0tyUJhZkoPCzJIcFGaW5KAwsyQHhZklOSjMLMlBYWZJDgozS3JQmFmSg8LMknL9mrmkjwO/AM4F3gT+MSLer9LuTWAfcBgYi4iL8vRrZs2Vd0ZxG/BURCwCnqps13JFRPytQ8Ks8+QNimuAn1de/xy4NufxzKwNKeL4C7tL2h0R8ydtvx8RH6vS7g3gfbLqYD+tVAOrdczJJQUvfPTRR4/7/NrV8PAwc+fWr2nRico6Lijv2K644ooXGpnlJ+9RSPof4PQqb/3wI5zP0ojYKWkh8KSk1yLimWoNXVKwc5V1XFDusTUiGRQR8Q+13pP0l/Gq5JLOAIZqHGNn5fuQpNVkFdCrBoWZtZ+89yjWAjdVXt8E/HJqA0knSZo3/hr4EvBKzn7NrInyBsUdwBcl/RH4YmUbSWdKGq8c9gngWUkvAX8AHo+IJ3L2a2ZNlGsdRUS8C3yhyv6dZLVGiYgdwOfy9GNmreWVmWaW5KAwsyQHhZklOSjMLMlBYWZJDgozS3JQmFmSg8LMkhwUZpbkoDCzJAeFmSU5KMwsyUFhZkkOCjNLclCYWZKDwsySHBRmluSgMLOkXEEh6auStko6IqlmbQBJyyX1S9ouqV41MTNrQ3lnFK8A11Hno/cldQF3A1cC5wM3Sjo/Z79m1kR5P1x3G4Ckes0uBrZXPmQXSY+QlSJ8NU/fZtY8uYKiQX3AW5O2B4BLajWeUlKQDRs2TOvJtcLw8LDH1WHKPLZG5CopGBHHFPypdogq+2oWPHVJwc5V1nFBucfWiFwlBRs0AJw9afssYGfOY5pZEzXj8ehzwCJJ50nqAW4gK0VoZh0i7+PRr0gaAC4DHpe0vrJ/oqRgRIwBtwLrgW3AoxGxNd9pm1kz5X3qsRpYXWX/REnByvY6YN3UdmbWGbwy08ySHBRmluSgMLMkB4WZJTkozCzJQWFmSQ4KM0tyUJhZkoPCzJIcFGaW5KAwsyQHhZklOSjMLMlBYWZJDgozS3JQmFmSg8LMkhwUZpbUrJKCb0raImmzpOfz9GlmzZe3ANB4ScGfNtD2ioh4J2d/ZtYCzSgpaGYdrln3KAL4taQXKiUDzayDNKOkIMDSiNgpaSHwpKTXIqJqBXTXHu1cZR0XlHtsjWhGScHxOh9ExJCk1WQVzqsGhWuPdq6yjgvKPbZGTPulh6STJM0bfw18iewmqJl1iGkvKQh8AnhW0kvAH4DHI+KJPP2aWXNNe0nBiNgBfC5PP2bWWl6ZaWZJDgozS3JQmFmSg8LMkhwUZpbkoDCzJAeFmSU5KMwsyUFhZkkOCjNLclCYWZKDwsySHBRmluSgMLMkB4WZJTkozCzJQWFmSQ4KM0tyUJhZUt4P110p6TVJL0taLWl+jXbLJfVL2i7ptjx9mlnz5Z1RPAl8JiL+BngduH1qA0ldwN3AlcD5wI2Szs/Zr5k1Ua6giIhfR8RYZXMjcFaVZhcD2yNiR0QcBB4BrsnTr5k1V95q5pN9E/hFlf19wFuTtgeAS2odZHJJQWBUUhmLBZ0GlLGye1nHBeUd2+JGGhVSe1TSD4Ex4KFqh6iyL2r1N7mkoKTnI+Ki1Dl2Go+r85R1bJKeb6Rd7tqjkm4Crga+EBHVAmAAOHvS9lnAzkZOzszaQ96nHsuBfwO+HBEHajR7Dlgk6TxJPcANwNo8/ZpZc+V96nEXMA94UtJmSffC0bVHKzc7bwXWA9uARyNia4PHvy/n+bUrj6vzlHVsDY1L1a8WzMw+5JWZZpbkoDCzpLYOikaXiHciSV+VtFXSEUkd/9itrMv0Jd0vaahs63kknS3pN5K2Vf4dfrde+7YOChpYIt7BXgGuA55p9YnkVfJl+g8Ay1t9EtNgDPh+RPw1cClwS72/s7YOigaXiHekiNgWEf2tPo+ClHaZfkQ8A7zX6vMoWkS8HREvVl7vI3si2VerfVsHxRTfBH7V6pOwqqot06/5j87ai6RzgSXA72u1KfJ3PY5LAUvE21YjYyuJj7RM39qHpLnAY8D3ImJvrXYtD4oCloi3rdTYSsTL9DuQpG6ykHgoIlbVa9vWlx4NLhG31vMy/Q4jScDPgG0R8ZNU+7YOCmosES8DSV+RNABcBjwuaX2rz+l45Vym39YkPQz8DlgsaUDSza0+p4IsBb4GfL7ys7VZ0lW1GnsJt5kltfuMwszagIPCzJIcFGaW5KAwsyQHhZklOSjMLMlBYWZJ/w/o7wtuSWNw6AAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# for i in range(len(res)):\n",
    "#     if i == 0:\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, aspect='equal', autoscale_on=False,\n",
    "                     xlim=(-2, 2), ylim=(-2, 2))\n",
    "ax.grid()\n",
    "\n",
    "line, = ax.plot([], [], 'o-', lw=2)\n",
    "# time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)\n",
    "# energy_text = ax.text(0.02, 0.90, '', transform=ax.transAxes)\n",
    "\n",
    "def init():\n",
    "    \"\"\"initialize animation\"\"\"\n",
    "    line.set_data([], [])\n",
    "    # time_text.set_text('')\n",
    "    # energy_text.set_text('')\n",
    "    return line,\n",
    "    # return line, time_text, energy_text\n",
    "\n",
    "class State:\n",
    "    def __init__(self, res):\n",
    "        self.i = 0\n",
    "        self.res = res\n",
    "    \n",
    "    def position(self):\n",
    "        x1 = np.cos(res[self.i, 0])\n",
    "        y1 = np.sin(res[self.i, 0])\n",
    "        x2 = np.cos(res[self.i, 0] + res[self.i, 1]) + x1\n",
    "        y2 = np.sin(res[self.i, 0] + res[self.i, 1]) + y1 \n",
    "        return (np.array([0, x1, x2]), np.array([0, y1, y2]))\n",
    "\n",
    "    def next(self):\n",
    "        self.i = self.i + 1\n",
    "\n",
    "state = State(res)\n",
    "print(state.position())\n",
    "def animate(i):\n",
    "    \"\"\"perform animation step\"\"\"\n",
    "    global state\n",
    "    # print(state.position())\n",
    "    line.set_data(*state.position())\n",
    "    state.next()\n",
    "    # time_text.set_text('time = %.1f' % pendulum.time_elapsed)\n",
    "    # energy_text.set_text('energy = %.3f J' % pendulum.energy())\n",
    "    return line,\n",
    "    # return line, time_text, energy_text\n",
    "\n",
    "# animate(0)\n",
    "ani = animation.FuncAnimation(fig, animate, blit=True, init_func=init)\n",
    "ani.save('double_pendulum.gif', fps=60, writer='pillow')\n",
    "plt.show()\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "b3ba2566441a7c06988d0923437866b63cedc61552a5af99d1f4fb67d367b25f"
  },
  "kernelspec": {
   "display_name": "Python 3.8.3 64-bit ('base': conda)",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
